import os
import argparse
import numpy as np
import torch
import pandas as pd
from model import SimpleRNAModel  # Assuming your model is in model.py

class Config:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    seq_vocab = "AUCG"
    coord_dims = 7 * 3  # Should match your training config

def predict_sequence(model, coord):
    """使用简化模型预测RNA序列"""
    # Preprocess coordinates
    coord = np.nan_to_num(coord)  # Handle NaN values
    coord_tensor = torch.tensor(coord.reshape(1, -1, Config.coord_dims), 
                              dtype=torch.float32).to(Config.device)
    length = torch.tensor([coord.shape[0]])
    
    # Predict
    with torch.no_grad():
        logits = model(coord_tensor, length)
        preds = logits[0].argmax(dim=1)  # Get predictions for first (and only) sequence
    
    return "".join([Config.seq_vocab[i] for i in preds.cpu().numpy()])

def predict(args):
    # 初始化模型
    model = SimpleRNAModel().to(Config.device)
    
    # 加载模型权重
    state_dict = torch.load("best_simple_model.pth", 
                          map_location=Config.device)
    model.load_state_dict(state_dict)
    model.eval()
    
    # 处理所有输入文件
    predictions = []
    for fname in os.listdir(args.input_coords):
        try:
            # 加载和处理坐标数据
            coord = np.load(os.path.join(args.input_coords, fname))
            
            # 获取PDB ID
            seq_id = os.path.splitext(fname)[0]
            
            # 使用模型预测序列
            pred_seq = predict_sequence(model, coord)
            
            predictions.append({
                "pdb_id": seq_id, 
                "seq": pred_seq
            })
            
        except Exception as e:
            print(f"处理文件 {fname} 时出错: {str(e)}")
            continue
    
    # 保存结果
    pd.DataFrame(predictions).to_csv(
        args.output,
        index=False,
        columns=["pdb_id", "seq"]
    )
    print(f"预测结果已保存到 {args.output}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="RNA序列预测")
    parser.add_argument("--input_coords", type=str, required=True,
                      help="输入坐标数据目录路径")
    parser.add_argument("--input_seqs", type=str, required=True,
                      help="输入序列数据目录路径（未使用，仅为兼容性保留）")
    parser.add_argument("--output", type=str, required=True,
                      help="输出CSV文件路径")
    
    args = parser.parse_args()
    
    # 创建输出目录（如果不存在）
    os.makedirs(os.path.dirname(args.output), exist_ok=True)
    
    predict(args)